from langchain_core.messages import HumanMessage, AIMessage, ToolMessage, FunctionMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough

# ✅ 模拟工具函数
@tool
def get_current_weather(location: str) -> str:
    """Get the weather for a location."""
    return f"The weather in {location} is sunny and 25°C."

# ✅ 注册工具到模型
llm = ChatOpenAI(model="gpt-4", temperature=0).bind_tools([get_current_weather], tool_choice="get_current_weather")

# ✅ 提示模板，带多轮对话历史
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful weather assistant."),
    MessagesPlaceholder("history"),
    ("human", "{input}")
])

# ✅ 输出结构化
parser = StrOutputParser()

# ✅ 构建链
chain = prompt | llm | parser

# ✅ 用户提问（第一轮对话）
user_input = "What's the weather in Tokyo?"
history = []

# 第一步：模型返回的是函数调用请求
intermediate_response = chain.invoke({
    "input": user_input,
    "history": history
})

print("\n🔧 模型返回（函数调用）:\n", intermediate_response)

# 模拟解析出来调用了哪个函数和参数（通常是函数名和参数）
# 实际情况中 intermediate_response 是 AIMessage，里面有 `tool_calls` 可提取
# 这里我们模拟调用结果
tool_result = get_current_weather.invoke({"location": "Tokyo"})  # 得到工具结果

# 第二步：添加 ToolMessage + 用户输入 + 模型响应
history.extend([
    HumanMessage(content=user_input),
    AIMessage(content="", additional_kwargs={
        "tool_calls": [{"id": "call-1", "function": {"name": "get_current_weather", "arguments": '{"location": "Tokyo"}'}}]
    }),
    ToolMessage(tool_call_id="call-1", content=tool_result)
])

# 第三步：继续调用模型，生成最终自然语言答复
final_response = chain.invoke({
    "input": "Okay, thanks!",
    "history": history
})

print("\n💬 模型最终自然语言回应:\n", final_response)
